{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Work in progress!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# III. Solving optimal control problems with DDP\n",
    "\n",
    "## III.a ABA dynamics for reaching a goal with Talos arm\n",
    "\n",
    "Our optimal control solver interacts with a defined ShootingProblem. A shooting problem represents a stack of action models in which an action model defines a specific node along the OC problem.\n",
    "\n",
    "First we need to create an action model from DifferentialFwdDynamics. We use it for building terminal and running action models. In this example, we employ an simpletic Euler integration rule.\n",
    "\n",
    "Next we define the set of cost functions for this problem. For this particular example, we formulate three running-cost functions:\n",
    "\n",
    "    goal-tracking cost, 𝑙𝑜𝑔(𝑓𝑋𝑑𝑜𝑜𝑋𝑓)\n",
    "\n",
    "state and control regularization; and ‖𝐱−𝐱𝑟𝑒𝑓‖,‖𝐮‖\n",
    "\n",
    "one terminal-cost:\n",
    "\n",
    "    goal cost. ‖𝐮𝑇‖\n",
    "\n",
    "First, let's create the common cost functions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import crocoddyl\n",
    "import pinocchio\n",
    "import numpy as np\n",
    "\n",
    "class DifferentialFreeFwdDynamicsModelDerived(\n",
    "    crocoddyl.DifferentialActionModelAbstract\n",
    "):\n",
    "    def __init__(self, state, actuationModel, costModel):\n",
    "        crocoddyl.DifferentialActionModelAbstract.__init__(\n",
    "            self, state, actuationModel.nu, costModel.nr\n",
    "        )\n",
    "        self.actuation = actuationModel\n",
    "        self.costs = costModel\n",
    "        self.enable_force = True\n",
    "        self.armature = np.matrix(np.zeros(0))\n",
    "\n",
    "    def calc(self, data, x, u=None):\n",
    "        if u is None:\n",
    "            q, v = x[: self.state.nq], x[-self.state.nv :]\n",
    "            pinocchio.computeAllTerms(self.state.pinocchio, data.pinocchio, q, v)\n",
    "            self.costs.calc(data.costs, x)\n",
    "            data.cost = data.costs.cost\n",
    "        else:\n",
    "            q, v = x[: self.state.nq], x[-self.state.nv :]\n",
    "            self.actuation.calc(data.actuation, x, u)\n",
    "            tau = data.actuation.tau\n",
    "            # Computing the dynamics using ABA or manually for armature case\n",
    "            if self.enable_force:\n",
    "                data.xout[:] = pinocchio.aba(\n",
    "                    self.state.pinocchio, data.pinocchio, q, v, tau\n",
    "                )\n",
    "            else:\n",
    "                pinocchio.computeAllTerms(self.state.pinocchio, data.pinocchio, q, v)\n",
    "                data.M = data.pinocchio.M\n",
    "                if self.armature.size == self.state.nv:\n",
    "                    data.M[range(self.state.nv), range(self.state.nv)] += self.armature\n",
    "                data.Minv = np.linalg.inv(data.M)\n",
    "                data.xout[:] = np.dot(data.Minv, (tau - data.pinocchio.nle))\n",
    "            # Computing the cost value and residuals\n",
    "            pinocchio.forwardKinematics(self.state.pinocchio, data.pinocchio, q, v)\n",
    "            pinocchio.updateFramePlacements(self.state.pinocchio, data.pinocchio)\n",
    "            self.costs.calc(data.costs, x, u)\n",
    "            data.cost = data.costs.cost\n",
    "\n",
    "    def calcDiff(self, data, x, u=None):\n",
    "        if u is None:\n",
    "            self.costs.calcDiff(data.costs, x)\n",
    "        else:\n",
    "            nq, nv = self.state.nq, self.state.nv\n",
    "            q, v = x[:nq], x[-nv:]\n",
    "            # Computing the actuation derivatives\n",
    "            self.actuation.calcDiff(data.actuation, x, u)\n",
    "            tau = data.actuation.tau\n",
    "            # Computing the dynamics derivatives\n",
    "            if self.enable_force:\n",
    "                pinocchio.computeABADerivatives(\n",
    "                    self.state.pinocchio, data.pinocchio, q, v, tau\n",
    "                )\n",
    "                ddq_dq = data.pinocchio.ddq_dq\n",
    "                ddq_dv = data.pinocchio.ddq_dv\n",
    "                data.Fx[:, :] = np.hstack([ddq_dq, ddq_dv]) + np.dot(\n",
    "                    data.pinocchio.Minv, data.actuation.dtau_dx\n",
    "                )\n",
    "                data.Fu[:, :] = np.dot(data.pinocchio.Minv, data.actuation.dtau_du)\n",
    "            else:\n",
    "                pinocchio.computeRNEADerivatives(\n",
    "                    self.state.pinocchio, data.pinocchio, q, v, data.xout\n",
    "                )\n",
    "                ddq_dq = np.dot(\n",
    "                    data.Minv, (data.actuation.dtau_dx[:, :nv] - data.pinocchio.dtau_dq)\n",
    "                )\n",
    "                ddq_dv = np.dot(\n",
    "                    data.Minv, (data.actuation.dtau_dx[:, nv:] - data.pinocchio.dtau_dv)\n",
    "                )\n",
    "                data.Fx[:, :] = np.hstack([ddq_dq, ddq_dv])\n",
    "                data.Fu[:, :] = np.dot(data.Minv, data.actuation.dtau_du)\n",
    "            # Computing the cost derivatives\n",
    "            self.costs.calcDiff(data.costs, x, u)\n",
    "\n",
    "    def createData(self):\n",
    "        data = DifferentialFreeFwdDynamicsDataDerived(self)\n",
    "        return data\n",
    "\n",
    "    def set_armature(self, armature):\n",
    "        if armature.size is not self.state.nv:\n",
    "            print(\"The armature dimension is wrong, we cannot set it.\")\n",
    "        else:\n",
    "            self.enable_force = False\n",
    "            self.armature = armature.T\n",
    "\n",
    "\n",
    "class DifferentialFreeFwdDynamicsDataDerived(crocoddyl.DifferentialActionDataAbstract):\n",
    "    def __init__(self, model):\n",
    "        crocoddyl.DifferentialActionDataAbstract.__init__(self, model)\n",
    "        self.pinocchio = pinocchio.Model.createData(model.state.pinocchio)\n",
    "        self.multibody = crocoddyl.DataCollectorMultibody(self.pinocchio)\n",
    "        self.actuation = model.actuation.createData()\n",
    "        self.costs = model.costs.createData(self.multibody)\n",
    "        self.costs.shareMemory(self)\n",
    "        self.Minv = None"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Building a differential action model for robot forward dynamics\n",
    "#### Loading the robot\n",
    "\n",
    "Crocoddyl offers several robot models for benchmarking our optimal control solvers (e.g. manipulators, humanoids, quadrupeds, etc). The collection of Talos models can be downloaded in Ubuntu with the APT package *robotpkg-talos-data*.\n",
    "\n",
    "Let's load a single Talos arm (left one):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import example_robot_data\n",
    "\n",
    "talos_arm = example_robot_data.load(\"talos_arm\")\n",
    "robot_model = talos_arm.model  # getting the Pinocchio model\n",
    "\n",
    "# Defining a initial state\n",
    "q0 = np.array([0.173046, 1.0, -0.52366, 0.0, 0.0, 0.1, -0.005])\n",
    "x0 = np.concatenate([q0, np.zeros(talos_arm.model.nv)])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create the cost functions\n",
    "target = np.array([0.4, 0.0, 0.4])\n",
    "state = crocoddyl.StateMultibody(robot_model)\n",
    "frameTranslationResidual = crocoddyl.ResidualModelFrameTranslation(\n",
    "    state, robot_model.getFrameId(\"gripper_left_joint\"), target\n",
    ")\n",
    "goalTrackingCost = crocoddyl.CostModelResidual(state, frameTranslationResidual)\n",
    "xRegCost = crocoddyl.CostModelResidual(state, crocoddyl.ResidualModelState(state))\n",
    "uRegCost = crocoddyl.CostModelResidual(state, crocoddyl.ResidualModelControl(state))\n",
    "\n",
    "# Create cost model per each action model\n",
    "runningCostModel = crocoddyl.CostModelSum(state)\n",
    "terminalCostModel = crocoddyl.CostModelSum(state)\n",
    "\n",
    "# Then let's added the running and terminal cost functions\n",
    "runningCostModel.addCost(\"gripperPose\", goalTrackingCost, 1e2)\n",
    "runningCostModel.addCost(\"stateReg\", xRegCost, 1e-4)\n",
    "runningCostModel.addCost(\"ctrlReg\", uRegCost, 1e-7)\n",
    "terminalCostModel.addCost(\"gripperPose\", goalTrackingCost, 1e5)\n",
    "terminalCostModel.addCost(\"stateReg\", xRegCost, 1e-4)\n",
    "terminalCostModel.addCost(\"ctrlReg\", uRegCost, 1e-7)\n",
    "\n",
    "# Running and terminal action models\n",
    "DT = 1e-3\n",
    "actuationModel = crocoddyl.ActuationModelFull(state)\n",
    "runningModel = crocoddyl.IntegratedActionModelEuler(\n",
    "    crocoddyl.DifferentialActionModelFreeFwdDynamics(\n",
    "        state, actuationModel, runningCostModel\n",
    "    ),\n",
    "    DT,\n",
    ")\n",
    "terminalModel = crocoddyl.IntegratedActionModelEuler(\n",
    "    crocoddyl.DifferentialActionModelFreeFwdDynamics(\n",
    "        state, actuationModel, terminalCostModel\n",
    "    ),\n",
    "    0.0,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Onces we have defined our shooting problem, we create a DDP solver object and pass some callback functions for analysing  its performance.\n",
    "\n",
    "Please note that:\n",
    "- CallbackDDPLogger: store the solution information.\n",
    "- CallbackDDPVerbose(level): printing message during the iterates.\n",
    "- CallbackDisplay(robot,rate): display the state trajectory using Gepetto viewer."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# For this optimal control problem, we define 250 knots (or running action\n",
    "# models) plus a terminal knot\n",
    "T = 250\n",
    "problem = crocoddyl.ShootingProblem(x0, [runningModel] * T, terminalModel)\n",
    "\n",
    "# Creating the DDP solver for this OC problem, defining a logger\n",
    "solver = crocoddyl.SolverFDDP(problem)\n",
    "log = crocoddyl.CallbackLogger()\n",
    "\n",
    "# Using the meshcat displayer, you could enable gepetto viewer for nicer view\n",
    "# display = crocoddyl.GepettoDisplay(talos_arm, 4, 4)\n",
    "display = crocoddyl.MeshcatDisplay(talos_arm, 4, 4)\n",
    "solver.setCallbacks([log, crocoddyl.CallbackVerbose(), crocoddyl.CallbackDisplay(display)])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Emdebbed meshcat in this cell\n",
    "display.robot.viewer.jupyter_cell()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Solving it with the DDP algorithm\n",
    "solver.solve()\n",
    "\n",
    "# Printing the reached position\n",
    "frame_idx = talos_arm.model.getFrameId(\"gripper_left_joint\")\n",
    "xT = solver.xs[-1]\n",
    "qT = xT[:talos_arm.model.nq]\n",
    "print()\n",
    "print(\"The reached pose by the wrist is\")\n",
    "print(talos_arm.framePlacement(qT, frame_idx))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "# # Plotting the solution and the DDP convergence\n",
    "crocoddyl.plotOCSolution(log.xs, log.us)\n",
    "crocoddyl.plotConvergence(\n",
    "    log.costs, log.pregs, log.dregs, log.grads, log.stops, log.steps\n",
    ")\n",
    "\n",
    "# Visualizing the solution in gepetto-viewer\n",
    "display.displayFromSolver(solver)"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
